# -*- coding: utf-8 -*-

"""图的顶点着色算法，可用于解决冲突场景，如考试科目的安排（哪些科目可以在同一时间考试），红绿灯调度（哪些车道可以同时通行）
   用冲突图来建模着色，比如在两门考试科目建立关系，看这两个科目是否有同一个同学选择，两个车道在路口有没有交叉...
   具体用例见经典的5路口红绿灯着色算法题。"""

class TrafficLight():
    def __init__(self) -> None:
        pass

    def coloring(self, crossing):
        """total:= 总共有的色号，用一个元组表示，最多每个顶点各一个颜色
           paint:= 已经着色的情况，用字典记录，{顶点：颜色}
           neighbor:= 临时变量，用于保存一个顶点周围顶点都有哪些颜色"""
        total = tuple(range(1, len(crossing)+1))
        paint = dict()

        def do(u):
            """深度优先的遍历"""
            neighbor = set()    # 统计顶点着了哪些色
            for v in crossing[u]:
                if w:= paint.get(v):
                    neighbor.add(w)

            for w in total:    # 从有序色号中取一个与周围色不同的最小色号
                if w not in neighbor:
                    current = w
                    break

            paint[u] = current    # 本顶点着色

            for v in crossing[u]:    # 深度优先，到下一个顶点
                if v not in paint.keys():
                    do(v)

        while True:    # 对非连通图的各支完成着色
            if left:= crossing.keys()-paint.keys():
                u = left.pop()
            else:
                break
            do(u)

        foo = dict()    # 屏蔽具体色号，将能着同样颜色的顶点汇集到一起
        for u, w in paint.items():
            if w in foo:
                foo[w].add(u)
            else:
                foo[w] = set([u])

        return tuple(foo.values())

if __name__ == '__main__':
    crossing = {
        'AB': {'BC', 'EA', 'BD', 'DA'},
        'AC': {'EA', 'BD', 'DA', 'DB', 'EB'},
        'AD': {'EA', 'EB', 'EC'},
        'BC': {'AB', 'EB', 'DB'},
        'BD': {'AB', 'AC', 'DA', 'EB', 'EC'},
        'BA': {},
        'DA': {'AB', 'AC', 'BD', 'EB', 'EC'},
        'DB': {'AC', 'BC', 'EC'},
        'DC': {},
        'EA': {'AB', 'AC', 'AD'},
        'EB': {'BC', 'BD', 'DA', 'AD', 'AC'},
        'EC': {'BD', 'DA', 'DB', 'AD'},
        'ED': {}
    }

    trafficlight = TrafficLight()
    light = trafficlight.coloring(crossing)
    print(light)
